import os
from dotenv import load_dotenv, find_dotenv 
_ = load_dotenv(find_dotenv())

from zhipuai import ZhipuAI
client = ZhipuAI()

def generate_text(prompt):
    messages = [{"role":"user", "content":prompt}]
    print(f"\n\nprompt=\n{prompt}")
    response = client.chat.completions.create(
        model="glm-4-plus",
        messages=messages
    )
    resultContent = response.choices[0].message.content.strip()
    print(f"resultContent=\n{resultContent}")
    return resultContent

import ast
# 思维树结点  
class ThoughtTreeNode:  
    def __init__(self, question, 
                 children=None, 
                 parentQuestion=None):  
        self.question = question  
        self.children = children if children is not None else [] 
        self.parentQuestion = parentQuestion 
        self.solution = None  
  
    def add_child(self, child_node):  
        self.children.append(child_node)  

    # 生成指定层级子话题
    def generate_children(self, level=1):
        level -= 1
        if self.parentQuestion is None:
            prompt = f"""讨论`{self.question}`这一话题，
                应当从哪几方面进行论论？请生成相关子话题列表。
                列表以`['子话题1','子话题2',...''子话题N'']`的形式返回，
                不要添加任何附加文字。""".replace(" ", "").replace("\t", "").replace("\n", "")
        else:
            prompt = f"""在父话题`{self.parentQuestion}`
                下讨论`{self.question}`这一话题，应当从哪几方面进行讨论？
                请生成相关子话题列表。
                列表以`['子话题1','子话题2',...''子话题N'']`的形式返回，
                不要添加任何附加文字。""".replace(" ", "").replace("\t", "").replace("\n", "")
        
        strChildren = generate_text(prompt)
        children = ast.literal_eval(strChildren) 
        for child in children:
            print(f"生成的{level}级子话题：{child}")
            if self.parentQuestion is None:
                child_parentQuestion = self.question
            else:
                child_parentQuestion = (
                    f"{self.parentQuestion}\\{self.question}"
                )
            child_node = ThoughtTreeNode(
                child, 
                parentQuestion=child_parentQuestion
            )

            if level > 0:
                child_node.generate_children(level)
            self.children.append(child_node) 
  
    # 生成解决方案
    def generate_solution(self):  
        if self.parentQuestion is None:
            prompt = f"""请给出`{self.question}`这一话题解决方案。"""
        else:
            prompt = f"""请给出`{self.question}`
                在父话题`{self.parentQuestion}`下的解决方案。"""

        if self.children is None:
            prompt = f"""{prompt}请详尽说明，1000字左右。"""
        else:
            prompt = f"""{prompt}概述即可，详细内容将由子话题描述，200字左右。"""

        response = generate_text(prompt)  
        self.solution = response  
          
        # 对子节点递归调用,生成对应话题解决方案  
        for child in self.children:  
            child.generate_solution()  
  
    # 输出思维树
    def print_tree(self, level=0):  
        print('  ' * level + str(self.question))  
        if self.solution:  
            print('  ' * level + '\n' + self.solution)  
        for child in self.children:  
            child.print_tree(level + 1)  
  
# 示例用法  
def main():  
    # 假设我们有一个问题  
    root_node = ThoughtTreeNode("如何学好人工智能编程？")  
  
    # 分解问题，向下分解两个层级，构建思维树
    root_node.generate_children(2) 
   
    # 生成解决方案  
    root_node.generate_solution()  
  
    # 打印思维树及其解决方案  
    root_node.print_tree()  
  
if __name__ == "__main__":  
    main()